'''
Created on May 2, 2011

@author: oabalbin
'''

import sys
import numpy as np
import array
from optparse import OptionParser
from collections import defaultdict,deque #, OrderedDict
from bx.intervals.intersection import Interval, IntervalTree

from exome.vcftools.vcf import *
import exome.vcftools.SNP_positions as snppos 

class Error(Exception):
    pass

class SnvError(Error):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class Gene:
    def __init__(self,start, end, value):
        self.start=start
        self.end=end
        self.value=value


class SequencesTree:
    
    def __init__(self):
        self.genome={}
        
        chrlist = ['chr1', 'chr2','chr3','chr4','chr5','chr6','chr7', 'chr8','chr9','chr10', 'chr11','chr12','chr6','chr7', 'chr8','chr9','chr10', 'chr11', 'chr12', \
                   'chr13', 'chr14','chr15','chr16','chr17','chr18','chr19', 'chr20','chr21','chr22', 'chrX','chrY']
        
        for chr in chrlist:
            self.genome[chr] = defaultdict(array.array('L'))


def actionable_genes(gene_list):
    '''
    '''
    actionable = IntervalTree()
    for gen in gene_list:
        start, end, obj = gen.start, gen.end, gen.value
        actionable.insert_interval( Interval(start, end, value=obj) )
    
    return actionable
        
        
def read_gene_list(ifile):
    '''
    Read a bed file with fields: chr, start, end, name, mark, strand
    '''
    ifile=open(ifile)
    actionable_genes=deque()
    for l in ifile:
        fields = l.strip('\n').split('\t')
        start, end, value=int(fields[1]),int(fields[2]), {'chr':fields[0], 'annot':fields[3],'strand':fields[5]}
        actionable_genes.append( Gene(start, end, value) )
    
    return actionable_genes
    
def mark_actionable_genes(ifile,snp_list):
    '''
    '''
    gene_list = read_gene_list(ifile)
    act_genes = actionable_genes(gene_list)
    act_snps = defaultdict()
    
    for snp in snp_list:
        loc = snp.split('@')
        chr,pos=loc[0],int(loc[1])
        genes = act_genes.find(pos,pos)
        for g in genes:
            if chr == g.value['chr']:
                act_snps[snp] = g.value['annot']
    
    return act_snps
    
    
def read_file(ivcf,pfname):
    
    thisvcf=VCF()
    thisvcf.readFromFile(ivcf)
    snvs = thisvcf.keys()
    #thisvcf.pickle(pfname)

    return snvs, thisvcf

def load_file(idump):
    thisvcf=VCF()
    thisvcf.readFromPickle(idump)
    
    return thisvcf
    

def read_file_vcf(ivcf):
    '''
    Reads a vcf returns a dictionary with snp location
    and set(id, ref and alt base)
    '''
    
    ivcf = open(ivcf)
    snps_dict = defaultdict(list)
    i=0
    for l in ivcf:
        i+=1
        if l.startswith('#'):
            continue

        fields = l.strip('\n').split('\t')
        loc = fields[0]+'@'+fields[1] 
        #loc = int(fields[1])
        snps_dict[loc].append([fields[0],fields[3],fields[4],fields[5]])
        
    print >> sys.stderr, "Reading vcf file is done"
    
    return snps_dict
    
def read_vcf_file_db(ivcf):
    '''
    '''
    
    snpsTree = SequencesTree()
    ivcf = open(ivcf)
    
    for i,l in enumerate(ivcf):
        
        if l.startswith('#'):
            continue
        
        fields = l.strip('\n').split('\t')
        #loc = fields[0]+'@'+fields[1] 
        chr, pos = fields[0], int(fields[1])
        snpsTree[chr].append(int(pos))
        
    print >> sys.stderr, "Reading vcf file is done"
    
    return snpsTree

    

def intersection_snvs(union_snvs, list_snvs_sets):
    '''
    
    '''
    # concesus matrix
    cm =  np.zeros( (len(union_snvs), len(list_snvs_sets)) )
    pm = defaultdict()
    
    for i, snv in enumerate(union_snvs):
        pm[snv]=i
        for j, thsnv_list in  enumerate(list_snvs_sets):
            if snv in thsnv_list: 
                cm[i,j]=1
    return pm, cm


def union_snvs_dict(union_snvs):
    '''
    '''
    pm = defaultdict()
    for i, snv in enumerate(union_snvs):
        pm[snv]=i
    return pm

def isec_snvs_annot(dict_snv_annot, dict_snvs_query):
    '''
    It intersect lists of snvs store 
    in a dictionary structure
    '''
    pm = defaultdict(dict)
    
    for ifile, snv_query in dict_snvs_query.iteritems():
        for annot, snv_annot in dict_snv_annot.iteritems():
            pm[ifile][annot] = set(snv_query).intersection(set(snv_annot))
            
    return pm

def isec_snvs_annot_union(snvs_union, snvs_annot_dbs):
    '''
    '''
    pm = defaultdict()
    
    for db, snv_db in snvs_annot_dbs.iteritems():
        pm[db] = set(snvs_union).intersection(set(snv_db))
            
    return pm


def base_type(base):
    '''
    determines the base type
    ns not standard base
    '''
    try:        
        snv_type= {'A':'pu','G':'pu','C':'pi','T':'pi'}[base]
        return snv_type
    except KeyError:
        return None


def titv_mutation(ref_base, alt_base):
    '''
    It determines if the mutation is transition
    or a transversion
    return ti for transition
    return tv for transversion
    '''
    transition=1
    transversion=0
    
    ref =base_type(ref_base)
    alt =base_type(alt_base)
    
    if (ref=='pu' and alt =='pu') or (ref=='pi'and alt=='pi'):
        return transition
    elif (ref=='pu' and alt =='pi') or (ref=='pi'and alt=='pu'):
        return transversion
 

def fields_table():
    '''
    '''
    '''
    format={"QUAL":'%10.2f',"DP":'%10.0f', "MQ":'%10.0f', "DP2":'%10.0f','DPREF':'%10.0f',
            'DPALT':'%10.0f',"AF":'%10.4f',"SB":'%10.4f'}
    '''
    fields={"QUAL":0,"DP":1, "MQ":2, 'DPREF':3,'DPALT':4,
            "AF":5,"SB":6,'VQSLOD':7,'HRun':8}
    
    format=['%10.2f','%10.0f','%10.0f','%10.0f',
            '%10.0f','%10.4f','%10.4f','%10.4f','%10.0f']

    names=["QUAL","DP", "MQ", 'DPREF','DPALT',
            "AF","SB", 'VQSLOD','HRun']
    
    gene_annot={'NON_SYNONYMOUS_CODING':1,'STOP_GAINED':2,'STOP_LOST':3,
                     'NON_SYNONYMOUS_CODING(SPLICE_SITE)':4,'STOP_GAINED(SPLICE_SITE)':5,
                     'ESSENTIAL_SPLICE_SITE':6, 'SPLICE_SITE':7,
                     'SYNONYMOUS_CODING':8,'3\'UTR':9,'5\'UTR':10,'INTRONIC':11,'INTERGENIC':12,
                     'DOWSTREAM':13,'UPSTREAM':14}

    return fields, format, names, gene_annot
    

def annotation_array(annotated_snvs, this_annot_snvs, nfiles):
    '''
    Annotate_snvs represent the snvs present in 
    particular vcf files thar are going to be queried. Dictuinary snc_postion: i
    This this_annot_snvs is double dictionary file_name,annotation database.
    '''
    nrow = len(annotated_snvs)
    ncol = len(this_annot_snvs)
    
    amat = np.zeros( (nrow, ncol) )
    j=0
    for f, snv_list in this_annot_snvs.iteritems():
        for snv in snv_list:
            i=annotated_snvs[snv]
            amat[i,j]=1
        j+=1
    
    return amat
        

def snv_concordance_to_database(query_snvs, ifile):
    '''
    common_snvs2databse: SNVs positions that are common with the database (set).
    snvs_in_database: ifile SNVs in the database in the format pos.base
    '''
    # read the pos.base file into two lists, the position list and the base list
    db_positions = snppos.read_position_base_file(ifile)
    # isec the commons_snvs2database and snvs_in_database
    isec_snvs = set(query_snvs.keys()).intersection(set(db_positions.keys()))
    # for the common snvs evaluate concordance.
    nrow = len(query_snvs)
    ncol = 1
    amat = np.zeros( (nrow,ncol) ) + 1
    
    for snv in isec_snvs:
        i = query_snvs[snv]
        db_base = db_positions[snv].split('>')[1]
        if db_base=='.':
            #print db_base, snv
            amat[i,0] = 0
#    print amat.shape
#    print len(amat[amat==1])
    return np.reshape(amat,(nrow,-1))



def build_concesus_table(annotated_snvs, dumped_files,ofile, mmapfile, 
                         complement, snvs_annot_dbs, hapmap_base=None):
    '''
    '''
    fields, format, names, gene_annot=fields_table()
    ncol = len(fields) # The one is for the gene_annotation column
    nrow = len(annotated_snvs)
    nfiles = len(dumped_files)
    mcol = len(dumped_files)+1
    ga = -1
    titvcol=2
    # Functional Annotation: Majority Vote and Present or not in HapMap
    
    print >> sys.stderr, "Building concensus table ...."
    full_annot = ['dbSNP','HapMap', 'Actionable', 'Validated']
    this_annot = set(full_annot).intersection(snvs_annot_dbs.keys())
    print this_annot
    
    amat = annotation_array(annotated_snvs, snvs_annot_dbs, nfiles) 
    # Matrix for transition - transversion snvs
    titvmat = np.zeros( (nrow,titvcol) )
        
    ############## Function
    #newpm=defaultdict()
    imat = np.zeros( (nrow,mcol) ) 
    gen_names=defaultdict()
    gen_annot_array=np.empty((nrow))
    
    k=0  
    for ft, df in dumped_files.iteritems(): # ft,df, ovcf
        ovcf=VCF()
        ovcf.readFromFile(df)
        # create a matrix for this file
        cmat = np.empty( (nrow,ncol) )
        
        for coord, i in annotated_snvs.iteritems():
            snv_list = ovcf[[coord]] # list of 1 snv 
            if not snv_list:
                cmat[i,:] = np.nan
                #print >> sys.stderr, "Not snps for this coordinate %s"%(coord)
            else:
                # Intersection matrix
                imat[i,k]=1
                
                for snv in snv_list:
                    # Filling annotation matrix
                    # print snv
                    loc = snv['COORD']
                    
                    if len(snv_list) > 1: 
                        print >> sys.stderr, "There is more than one snp with the same location"
                    
                    transition = titv_mutation(snv['REF'], snv['ALT'])
                    #print snv['REF'], snv['ALT'], transition
                    if transition:
                        titvmat[i,0]=1
                    else:
                        titvmat[i,1]=1
                       
                    # Properties of the snp
                    for prop,j in fields.iteritems():
                        
                        try:
                            cmat[i,j] = snv[prop]
                        except KeyError:
                            print >> sys.stderr, "%s, Error with this snv %s, it is missing field %s"%(ft,coord,prop)
                            if prop=="SB":
                                # This number means the field was not present in the vcf
                                cmat[i,j]=999
                            else:
                                cmat[i,j] = np.nan
                    
                    # The part below handles the snp annotation and the gene name associated to it. 
                    snv_annotation = set(gene_annot.keys()).intersection(snv.keys())
                    if snv_annotation:
                        tmpid=100                        
                        for gt in snv_annotation:
                            id = gene_annot[gt]
                            if id < tmpid:
                                tmp,tmpid=gt,id
                        gen_annot_array[i]=id
                        gen_names[i]=snv[tmp]
                        
                    else:
                        gen_annot_array[i]=999
                        gen_names[i]='nan'
     
        # concatenate matrix for this file to gral matrix
        if k==0:
            Cmat=np.copy(cmat)
        else:
            Cmat = np.hstack( (Cmat,cmat) )
        k+=1
        
    
    # Get the majority vote for the snv calls
    fformat=[]
    fnames=[]
    fmethods=[]
    for i in range(len(dumped_files)):
        tn = ['M%d_%s'%(i,n) for n in names]
        fformat+=format
        fnames+=tn
        fmethods.append('M%d'%(i)) 
    
    this_annot = list(this_annot)
    print fmethods
    this_annot.insert(0,'Majority')
    print this_annot
    fmethods=fmethods+this_annot
    fmethods.reverse()
    [fformat.insert(0,'%d') for i in range(imat.shape[1])]
    [fnames.insert(0,f) for f in fmethods]
    
    fnames.extend(['Ti','Tv'])
    # Concatenate the intersection and property matrices
    print imat.shape, mcol
    imat[:,nfiles]=np.sum(imat[:,:mcol], axis=1)
    # Concatenate imat and annotation
    iamat = np.hstack( (imat,amat) )
    # Concatenate the annotation matrix
    Cmat = np.hstack( (iamat,Cmat) )
    # Concatenate the matrix and the titv matrix
    Cmat = np.hstack( (Cmat,titvmat) )
    
    ### Calculate Genotype call accuracy.
    # Number of calls incorrect calls. Calls that were called as variant but 
    # they are really a reference position.
    print "Before gmat", Cmat.shape
    if hapmap_base is not None:
        gmat = snv_concordance_to_database(annotated_snvs, hapmap_base)
        print "The total number of incorrect calls are:", len(gmat[gmat==0])
        Cmat = np.hstack( (Cmat,gmat) )
        fnames.extend(['GI'])
        print "Calculated gmat", Cmat.shape, fnames
     
    
    write_table(Cmat,ofile,",".join(fformat),mmapfile, annotated_snvs, fnames, gen_annot_array, gen_names)
    
    return fnames


def write_table(mat,fname, format, mmapfile, snps, header, gen_annot=None, gene_names=None):
    '''
    np.memmap(array, dtype='float')
    numpy.savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n')
    Sequence formtas for sved file, so reduce size.
    1.0000          1.0000          1.0000        120.0000        207.0000         60.0000        225.0000        207.0000         60.0000       5594.8100    207.0000          60.0000
    Ask for the format sequence fromthe beggining as input parameters
    '''
    
    #np.savetxt(fname,mat,fmt=format,delimiter='\t')#,newline='\n') '%d,%d,%d,%d,%10.2f,%10.0f,%10.0f,%10.2f,%10.0f,%10.0f,%10.2f,%10.0f,%10.0f'
    cmat = np.memmap(mmapfile,dtype='float32',mode='w+',shape=(mat.shape[0],mat.shape[1]) )
    cmat[:] = mat[:]
    
    
    ofile = open(fname,'w')
    
    # print snvs list
    if gene_names is not None and gen_annot is not None:
        ofile.write("snpID\t"+",".join(header).replace(',','\t')+"\tmut_type"+"\tgene_name"+'\n')
        for s,i in snps.iteritems():
            l = [s] + map(str,list(mat[i,:])) + [str(gen_annot[i])] + ['\''+",".join(gene_names[i]).replace(',','&')+'\'']
            ofile.write(",".join(l).replace(',','\t')+'\n')
        ofile.close()
    else:
        ofile.write("snpID\t"+",".join(header).replace(',','\t')+'\n')
        for s,i in snps.iteritems():
            l = [s] + map(str,list(mat[i,:])) + ['\''+'nan'+'\''] + ['\''+'nan'+'\'']
            ofile.write(",".join(l).replace(',','\t')+'\n')
        ofile.close()

        
        

def main_concesus(vcf_files, list_annot_files, ofile, mmapfile='',complement=False):
    '''
    '''
    snvs_union = set()
    pfiles=defaultdict()
    snvs_dict = defaultdict()
    snvs_annot_dbs = defaultdict()
    
    
    for ft, f in vcf_files.iteritems():
        # Generate a positions files for query vcf.
        fpos = f.replace('.vcf','.pos')
        snppos.read_file(f,fpos)
        snvs_dict[f]=snppos.read_positions(fpos)
        pf = f+'.dump'
        #pfiles[ft]=pf
        print >> sys.stderr, "Reading vcf file %s"%(f)
        thsnvs,ovcf = read_file(f,pf)
        pfiles[ft]=f #ovcf
        snvs_union.update(set(thsnvs))
        
    
    cm=[]
    snvs_query = union_snvs_dict(snvs_union)
    # Actionable, dbSNP and HapMap SNPs
    snvs_annot_dbs['Actionable']=mark_actionable_genes(list_annot_files[0], snvs_union)
    snvs_annot_dbs['dbSNP'] = snppos.read_positions(list_annot_files[1])
    snvs_annot_dbs['HapMap'] = snppos.read_positions(list_annot_files[2])
    #snvs_annot_dbs['Validated'] = read_file_vcf(list_annot_files[3])
    
    # Intersection of snvs and annotated snvs
    #annotated_snvs = isec_snvs_annot(snvs_dict, snvs_annot_dbs)
    annotated_snvs = isec_snvs_annot_union(snvs_union, snvs_annot_dbs)
    
    print >> sys.stderr, "Building concensus"
    '''
    build_concesus_table(pm, cm, dumped_files,ofile, mmapfile, 
                         complement, act_snps=None, dbsnp=None, hapmap=None, validated=None)
    Annotate pm, cm, pfiles, ofile, mmapfile, complement,
    '''
    if len(list_annot_files)==4:
        print "Doing concordance"
        
        header = build_concesus_table(snvs_query, pfiles, ofile, mmapfile,
                                  complement, annotated_snvs,list_annot_files[3])
    else:
        header = build_concesus_table(snvs_query, pfiles, ofile, mmapfile,
                                      complement, annotated_snvs)
    
    hfile = outfile.replace('.txt','.header')
    # print table header
    hfile = open(hfile,'w')
    hfile.write(",".join(header).replace(',',',\t')+'\n')
    hfile.close()
    # print snvs list
    snvs = outfile.replace('.txt','.snvs')
    ofile = open(snvs,'w')
    for s,i in snvs_query.iteritems():
        ofile.write(s+'\t'+str(i)+'\n')
    ofile.close()
    print >> sys.stderr,"DONE!"
        
        
if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] ")
    
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    # Specific options
    optionparser.add_option("-f", "--vcf_files", action="append", type="str",
                            dest="vcf_files",
                            help="others_vcf_files. Pass vcf file in the following order: GATK, SAM, Other")
    optionparser.add_option("-a", "--annot_files", action="append", type="str",
                            dest="annot_files",
                            help="other annotation files. ")
    optionparser.add_option("-o", "--output_file", dest="output_file",
                            help="output_file")
    optionparser.add_option("-c", "--complement", action='store_true', default=False,
                            help="output_file")
    
    


    (options, args) = optionparser.parse_args()
    outfile =options.output_file+'.txt'
    mmapfile =options.output_file+'.memmap'
    print options
    if len(options.vcf_files) < 2:
        print >> sys.stderr, "You need a minimum of 2 vcf files to do the analysis"
        sys.exit(0)
    elif len(options.vcf_files) ==3:
        vcf_files = {'GATK':options.vcf_files[0],'SAM':options.vcf_files[1], 
                     'VARS':options.vcf_files[2]}
    elif len(options.vcf_files) ==2:
        vcf_files = {'GATK':options.vcf_files[0],'SAM':options.vcf_files[1]}
        
    
    main_concesus(vcf_files, options.annot_files, outfile, mmapfile, options.complement)